import pandas as pd
import numpy as np
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import warnings
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf

warnings.filterwarnings('ignore')  # 临时忽略警告，实际使用时可移除


class EcommerceSearchPredictor:
    """阿里电商搜索关键词预测模型"""

    def __init__(self):
        """初始化模型"""
        self.model = None
        self.model_fit = None
        self.data = None
        self.keyword = None
        self.freq = 'D'  # 明确设置时间频率为日

    def load_data(self, file_path=None, keyword=None, sample_data=True):
        """
        加载电商搜索数据

        参数:
            file_path: CSV文件路径，如果为None则使用示例数据
            keyword: 要分析的关键词
            sample_data: 是否使用示例数据
            基础搜索量在 1000-10000 之间随机波动
            每周有周期性模式，周末搜索量会增加
            包含三个主要促销活动（618、双 11、双 12），这些时段搜索量显著提升
            春节前也有小幅搜索量增长
            添加了随机噪声使数据更真实
        """
        if sample_data:
            # 生成示例数据
            dates = pd.date_range(start='2023-01-01', end='2024-12-31', freq=self.freq)
            np.random.seed(42)

            # 基础搜索量（增加长期增长趋势）
            base_volume = np.random.randint(1000, 10000, size=len(dates))
            trend = np.linspace(1, 1.3, len(dates))  # 2年10%增长趋势
            base_volume = (base_volume * trend).astype(int)

            # 添加季节性趋势
            weekly_pattern = [1.0, 1.0, 1.0, 1.0, 1.2, 1.5, 1.5]  # 周末搜索量增加
            seasonal_volume = np.array([base_volume[i] * weekly_pattern[date.weekday()]
                                        for i, date in enumerate(dates)])

            # 添加促销活动效应
            promotional_volume = seasonal_volume.copy()
            for i, date in enumerate(dates):
                # 618大促
                if date.month == 6 and 15 <= date.day <= 20:
                    promotional_volume[i] *= 3
                # 双11
                if date.month == 11 and 8 <= date.day <= 14:
                    promotional_volume[i] *= 5
                # 双12
                if date.month == 12 and 10 <= date.day <= 14:
                    promotional_volume[i] *= 2.5
                # 春节前
                if (date.month == 1 and 20 <= date.day <= 31) or (date.month == 2 and date.day <= 10):
                    promotional_volume[i] *= 1.8

            # 添加随机噪声
            noise = np.random.normal(0, 0.1, size=len(dates))
            final_volume = promotional_volume * (1 + noise)

            self.data = pd.DataFrame({
                'date': dates,
                'search_volume': final_volume.astype(int)
            })
            self.keyword = "智能手机" if keyword is None else keyword
        else:
            # 从CSV加载数据
            if file_path is None:
                raise ValueError("当sample_data为False时，必须提供file_path")
            self.data = pd.read_csv(file_path)
            self.keyword = keyword if keyword is not None else "未指定关键词"
            # 尝试从数据中推断频率
            self.freq = pd.infer_freq(self.data['date']) or 'D'

        # 设置日期索引并确保频率一致
        self.data.set_index('date', inplace=True)
        self.data = self.data.asfreq(self.freq)
        return self.data

    def check_stationarity(self):
        """检查时间序列平稳性"""
        result = adfuller(self.data['search_volume'].dropna())
        print(f'ADF检验统计量: {result[0]:.4f}')
        print(f'p值: {result[1]:.4f}')
        print('临界值:')
        for key, value in result[4].items():
            print(f'  {key}: {value:.4f}')

        if result[1] > 0.05:
            print("序列非平稳，可能需要差分")
        else:
            print("序列平稳")
        return result[1] > 0.05

    def suggest_arima_params(self):
        """通过ACF和PACF图建议ARIMA参数"""
        plt.figure(figsize=(12, 6))
        plt.subplot(211)
        plot_acf(self.data['search_volume'].dropna(), ax=plt.gca())
        plt.subplot(212)
        plot_pacf(self.data['search_volume'].dropna(), ax=plt.gca())
        plt.tight_layout()
        plt.show()

        print("请观察ACF和PACF图以确定p和q值:")
        print("- ACF图的拖尾情况建议q值")
        print("- PACF图的截尾情况建议p值")
        print("(示例: p=5, d=1, q=0)")
        return 5, 1, 0  # 默认返回示例参数

    def train_model(self, p=5, d=1, q=0):
        """
        训练ARIMA模型，增加自动参数选择功能

        参数:
            p, d, q: ARIMA模型参数
        """
        if self.data is None:
            self.load_data()

        # 检查平稳性以确定d值
        if d is None:
            d = 1 if self.check_stationarity() else 0

        # 训练ARIMA模型
        self.model = ARIMA(self.data['search_volume'], order=(p, d, q), freq=self.freq)
        self.model_fit = self.model.fit()
        print(f"模型训练完成 - AIC: {self.model_fit.aic:.2f}")
        return self.model_fit

    def predict(self, days=30, plot=True):
        """
        预测未来搜索趋势，改进日期生成逻辑

        参数:
            days: 预测天数
            plot: 是否绘制预测图表
        """
        if self.model_fit is None:
            self.train_model()

        # 预测未来值
        forecast = self.model_fit.forecast(steps=days)

        # 创建预测日期（从数据最后一天的下一天开始）
        last_date = self.data.index[-1]
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=days, freq=self.freq)

        # 创建预测数据框
        forecast_df = pd.DataFrame({
            'date': forecast_dates,
            'predicted_search_volume': forecast
        })
        forecast_df.set_index('date', inplace=True)

        if plot:
            # 绘制历史数据和预测数据（优化中文显示）
            plt.figure(figsize=(14, 7))
            plt.plot(self.data.index, self.data['search_volume'], label='历史搜索量')
            plt.plot(forecast_df.index, forecast_df['predicted_search_volume'],
                     label='预测搜索量', color='red', linestyle='--')

            # 标记重要促销日期
            promotion_dates = [
                pd.Timestamp('2023-06-18'), pd.Timestamp('2023-11-11'), pd.Timestamp('2023-12-12'),
                pd.Timestamp('2024-06-18'), pd.Timestamp('2024-11-11'), pd.Timestamp('2024-12-12')
            ]
            for date in promotion_dates:
                if date in self.data.index:
                    plt.axvline(x=date, color='green', linestyle='--', alpha=0.5)

            plt.title(f'{self.keyword}搜索量预测趋势', fontproperties='SimHei')
            plt.xlabel('日期', fontproperties='SimHei')
            plt.ylabel('搜索量', fontproperties='SimHei')
            plt.legend(prop={'family': 'SimHei'})
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.show()

        return forecast_df

    def evaluate(self, test_size=30):
        """
        评估模型性能，使用更合理的测试集划分

        参数:
            test_size: 用于测试的数据点数量
        """
        if self.data is None:
            self.load_data()

        # 确保有足够的数据用于测试
        if len(self.data) < test_size * 2:
            test_size = len(self.data) // 2
            print(f"数据量不足，调整测试集大小为: {test_size}")

        # 分割数据（使用时间序列的前80%作为训练集）
        train_size = len(self.data) - test_size
        train_data = self.data.iloc[:train_size]
        test_data = self.data.iloc[train_size:]

        # 训练模型
        model = ARIMA(train_data['search_volume'], order=(5, 1, 0), freq=self.freq)
        model_fit = model.fit()

        # 预测
        predictions = model_fit.forecast(steps=test_size)

        # 计算误差
        mse = mean_squared_error(test_data['search_volume'], predictions)
        rmse = np.sqrt(mse)
        # 避免除以零
        mask = test_data['search_volume'] > 0
        mape = np.mean(np.abs((test_data['search_volume'][mask] - predictions[mask]) /
                              test_data['search_volume'][mask])) * 100

        print(f"评估结果:")
        print(f"  训练集大小: {len(train_data)}")
        print(f"  测试集大小: {len(test_data)}")
        print(f"  MSE: {mse:.2f}")
        print(f"  RMSE: {rmse:.2f}")
        print(f"  MAPE: {mape:.2f}%")

        return {
            'mse': mse,
            'rmse': rmse,
            'mape': mape
        }


# 使用示例
if __name__ == "__main__":
    # 解决中文字体问题
    plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
    plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

    # 创建预测器实例
    predictor = EcommerceSearchPredictor()

    # 加载示例数据
    predictor.load_data(keyword="智能手机")

    # 检查序列平稳性并建议参数
    # predictor.check_stationarity()
    # p, d, q = predictor.suggest_arima_params()

    # 训练模型（使用默认参数或取消注释上面两行使用建议参数）
    predictor.train_model()

    # 预测未来90天
    forecast = predictor.predict(days=90)

    # 评估模型
    metrics = predictor.evaluate()

    print("\n未来3天预测结果:")
    print(forecast.head(3))